import cv2
import numpy as np

# 读取图像
image = cv2.imread('input_data/src2.jpg')

# 转换为灰度图像
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# 使用 CLAHE 增强局部对比度
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
equalized_image = clahe.apply(gray_image)

# 应用形态学操作（闭运算：填充小孔，开运算：去除小噪声）
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))  # 定义形态学核
filtered_image = cv2.morphologyEx(gray_image, cv2.MORPH_CLOSE, kernel)  # 闭运算
filtered_image = cv2.morphologyEx(filtered_image, cv2.MORPH_OPEN, kernel)  # 开运算

# 应用高斯模糊，减少噪声并平滑图像 ==> 效果不大好
# blurred_image = cv2.GaussianBlur(equalized_image, (3,3), 0)  # 高斯核大小为 (5,5)

# 应用多级阈值分割，形成阶梯状灰度
# 假设我们希望将灰度分为10个等级
levels = 10
max_gray = 255
increment = max_gray // levels

# 创建阶梯状的灰度分布
for i in range(levels):
    lower = i * increment
    upper = (i+1) * increment
    filtered_image[(filtered_image >= lower) & (filtered_image < upper)] = lower + (upper - lower) // 2

# 将处理后的灰度图像转换为伪彩色图像，使用不同的颜色映射
heatmap = cv2.applyColorMap(filtered_image, cv2.COLORMAP_BONE)

# 调整热图的亮度，使亮面差异更明显
# beta小于0，降低亮度
heatmap = cv2.convertScaleAbs(heatmap, alpha=1, beta=-50)  

# 定义透明度
alpha = 0.5

# 将热图叠加到原图上
result = cv2.addWeighted(image, 1, heatmap, alpha, 0)

# 显示结果
cv2.imshow('Original Image', image)
# cv2.imshow('CLAHE Equalized Image', equalized_image)
# cv2.imshow('Filtered Image', filtered_image)
cv2.imshow('Heatmap', heatmap)
# cv2.imshow('Result Image', result)

# 保存处理后的图像
# cv2.imwrite('output_data/heatmap.jpg', heatmap)
# cv2.imwrite('output_data/result.jpg', result)

# 等待按键按下后关闭所有窗口
cv2.waitKey(0)
cv2.destroyAllWindows()